import os
import sys
import numpy as np
import pandas as pd
import paddle
import matplotlib
matplotlib.use('Agg')  # 避免 GUI 环境问题
import matplotlib.pyplot as plt

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

# 导入模块
from models.time_series_data import TimeSeriesData
from models.efficient_time_llm import EfficientTimeLLM
from models.train_eval import train_model, evaluate_model
from scripts.predict_ensemble import sliding_window_predict


# 创建保存目录
os.makedirs('results', exist_ok=True)
os.makedirs('models', exist_ok=True)


def load_and_preprocess_data(data_path, target_metrics):
    """
    自定义数据加载与预处理函数
    动态识别 metric_name 并转换为宽格式
    """
    # 1. 读取原始数据
    df = pd.read_csv(data_path)

    # 2. 检查是否为长格式数据
    if 'metric_name' in df.columns and 'value' in df.columns and 'timestamp_value' in df.columns:
        # 3. 动态获取所有 metric_name
        available_metrics = df['metric_name'].unique().tolist()
        print(f"🔍 自动检测到以下指标: {available_metrics}")

        # 4. 过滤掉 CSV 中不存在的 target_metrics
        valid_metrics = [m for m in target_metrics if m in available_metrics]
        if not valid_metrics:
            raise ValueError("没有在 CSV 中找到任何目标指标，请检查 target_metrics 设置")

        # 5. 转换为宽格式
        df_wide = df.pivot(index='timestamp_value', columns='metric_name', values='value').reset_index()

        return df_wide, valid_metrics

    else:
        raise ValueError("CSV 文件格式不支持，请确保包含 'metric_name', 'value', 'timestamp_value' 列")


if __name__ == "__main__":
    # 1. 数据加载
    data_path = "data/cpu_test.csv"
    if not os.path.exists(data_path):
        raise FileNotFoundError(f"找不到数据文件: {data_path}")

    # 自定义目标指标
    target_metrics = ["usage_active", "mem_used", "disk_io", "network_in", "network_out"]

    print("加载并预处理数据...")
    df_wide, detected_metrics = load_and_preprocess_data(data_path, target_metrics)

    # 保存为临时文件（供 TimeSeriesData 使用）
    temp_data_path = "data/temp_preprocessed.csv"
    df_wide.to_csv(temp_data_path, index=False)

    print(f"已生成预处理后的宽格式文件: {temp_data_path}")

    # 2. 初始化 TimeSeriesData
    print("初始化数据处理器...")
    data_loader = TimeSeriesData(
        temp_data_path,
        target_metrics=detected_metrics  # 使用检测到的有效指标
    )
    X, time_X, y = data_loader.create_dataset()

    print(f"数据集大小: X={X.shape}, time_X={time_X.shape}, y={y.shape}")

    # 3. 数据集划分
    split_idx = int(len(X) * 0.8)
    X_train, X_test = X[:split_idx], X[split_idx:]
    time_train, time_test = time_X[:split_idx], time_X[split_idx:]
    y_train, y_test = y[:split_idx], y[split_idx:]

    print(f"训练集: {len(X_train)} 样本, 测试集: {len(X_test)} 样本")

    # 4. 初始化模型
    print("初始化模型...")
    model = EfficientTimeLLM(
        seq_len=96,
        pred_len=24,
        d_model=128,
        num_layers=3,
        top_k=5,
        n_vars=X_train.shape[2]  # 自动获取多变量数量
    )

    # 自动使用 GPU（如果可用）
    use_gpu = paddle.is_compiled_with_cuda()
    if use_gpu:
        model = model.cuda()
        print("⚡ 使用 GPU 进行训练")

    # 5. 训练主模型
    print("🏋开始训练主模型...")
    losses = train_model(model, X_train, time_train, y_train, epochs=100, batch_size=32)

    # 保存主模型
    paddle.save(model.state_dict(), 'models/main_model.pdparams')
    print("主模型已保存")

    # 6. 绘制训练损失曲线
    print("绘制训练损失曲线...")
    plt.figure(figsize=(10, 5))
    plt.plot(losses, label='Training Loss')
    plt.title('Training Loss Curve')
    plt.xlabel('Epoch')
    plt.ylabel('Loss')
    plt.grid(True)
    plt.legend()
    plt.savefig('results/training_loss.png')
    plt.close()

    # 7. 评估主模型
    print("评估主模型...")
    test_loss, y_pred = evaluate_model(model, X_test, time_test, y_test, data_loader.scaler, return_predictions=True)
    print(f"测试集 MSE: {test_loss:.6f}")

    # 8. 绘制测试集预测对比图
    print("绘制测试集预测对比图...")
    plt.figure(figsize=(12, 6))
    plt.plot(y_test[:100], label='True')
    plt.plot(y_pred[:100], label='Predicted')
    plt.title('Test Set Prediction vs True')
    plt.xlabel('Sample')
    plt.ylabel('Value')
    plt.grid(True)
    plt.legend()
    plt.savefig('results/test_prediction.png')
    plt.close()

    # 9. 滑动预测
    print("进行未来预测...")
    latest_data = X_test[-1]  # 最新数据
    latest_time = time_test[-1]

    future_pred = sliding_window_predict(model, latest_data, latest_time, pred_steps=24)
    future_pred = data_loader.scaler.inverse_transform(np.array(future_pred).reshape(-1, 1)).flatten()

    print("未来24分钟预测值:")
    print(future_pred)

    # 10. 绘制未来24分钟预测图
    print("生成未来预测图...")
    plt.figure(figsize=(10, 5))
    plt.plot(future_pred, marker='o', label='Predicted')
    plt.title('Future 24-Minute Prediction')
    plt.xlabel('Minutes Ahead')
    plt.ylabel('Value')
    plt.grid(True)
    plt.legend()
    plt.savefig('results/future_prediction.png')
    plt.close()

    # 11. 绘制历史 + 预测拼接图
    print("生成历史+预测拼接图...")
    last_96 = data_loader.scaler.inverse_transform(X_test[-1][:, 0].reshape(-1, 1)).flatten()
    combined = np.concatenate([last_96, future_pred], axis=0)

    plt.figure(figsize=(12, 6))
    plt.plot(combined, marker='o', label='History + Prediction')
    plt.axvline(x=96, color='r', linestyle='--', label='Prediction Start')
    plt.title('Historical + Future 24-Minute Prediction')
    plt.xlabel('Time (minutes)')
    plt.ylabel('Value')
    plt.grid(True)
    plt.legend()
    plt.savefig('results/history_prediction.png')
    plt.close()

    print("所有结果已保存为图像文件")
